package dd.lo.classifier;

import dd.lo.util.ImageUtils;
import javafx.application.Application;
import javafx.application.Platform;
import javafx.beans.property.SimpleStringProperty;
import javafx.beans.property.StringProperty;
import javafx.concurrent.Task;
import javafx.scene.Scene;
import javafx.scene.control.Button;
import javafx.scene.control.TextField;
import javafx.scene.image.Image;
import javafx.scene.image.ImageView;
import javafx.scene.layout.BorderPane;
import javafx.scene.layout.FlowPane;
import javafx.scene.text.Text;
import javafx.stage.DirectoryChooser;
import javafx.stage.Stage;
import org.apache.commons.lang3.math.NumberUtils;
import org.opencv.core.Mat;
import org.opencv.core.MatOfByte;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;

public class MnistViewer extends Application {

    private static final int DEFAULT_WIDTH = 1080;
    private static final int DEFAULT_HEIGHT = 768;

    public static void start(String[] args) {
        launch(args);
    }

    private StringProperty pic;

    private ImageView imageView;

    private Text label;

    private MnistReader trainImgReader;
    private MnistReader trainLabelReader;
    private MnistReader testImgReader;
    private MnistReader testLabelReader;

    private void readMnist(File directory) throws IOException {
        String trainLabelFilePath = directory.getAbsolutePath() + "/train-labels-idx1-ubyte.gz";
        File trainLabelFile = new File(trainLabelFilePath);
        if (!trainLabelFile.exists()) throw new FileNotFoundException(String.format("训练集标签文件未能找到（%s）", trainLabelFilePath));
        trainLabelReader = new MnistReader(trainLabelFile.getAbsolutePath(), MnistReader.TYPE_LABEL);
        String trainImgFilePath = directory.getAbsolutePath() + "/train-images-idx3-ubyte.gz";
        File trainImgFile = new File(trainImgFilePath);
        if (!trainImgFile.exists()) throw new FileNotFoundException(String.format("训练集图片文件未能找到（%s）", trainImgFilePath));
        trainImgReader = new MnistReader(trainImgFile.getAbsolutePath(), MnistReader.TYPE_IMAGE);
    }

    private void showPic(int i) {
        byte[] labelData = trainLabelReader.get(i);
        int l = labelData[0] & 0xFF;
        byte[] imageData = trainImgReader.get(i);
        int height = trainImgReader.getHeight(), width = trainImgReader.getWidth();
        Task<Void> task = new Task<Void>() {
            @Override
            protected Void call() throws IOException {
                BufferedImage image = new BufferedImage(trainImgReader.getWidth(), trainImgReader.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
                int idx = 0;
                for (int y = 0; y < height; ++y) {
                    for (int x = 0; x < width; ++x) {
                        image.setRGB(x, y, imageData[idx] & 0xFF);
                        ++idx;
                    }
                }
                ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
                ImageIO.write(image, "jpeg", outputStream);
                Mat src = Imgcodecs.imdecode(new MatOfByte(outputStream.toByteArray()), Imgcodecs.IMREAD_UNCHANGED);
                Imgproc.threshold(src, src, 5, 255, Imgproc.THRESH_BINARY);
                //调整图片大小以用于展示
                ImageUtils.fitImgSize(src, DEFAULT_WIDTH, DEFAULT_HEIGHT);
                MatOfByte buffer = new MatOfByte();
                Imgcodecs.imencode(".jpeg", src, buffer);
                Platform.runLater(() -> {
                    imageView.setImage(new Image(new ByteArrayInputStream(buffer.toArray())));
                    label.setText(l + "");
                });
                return null;
            }
        };
        new Thread(task).start();
    }

    @Override
    public void start(Stage stage) {
        //Warning text
        Text warning = new Text();
        //Label text
        label = new Text();
        //Directory chooser
        Button dirChooser = new Button("选择文件夹");
        dirChooser.setOnMouseClicked(e -> {
            DirectoryChooser directoryChooser = new DirectoryChooser();
            directoryChooser.setTitle("请选择文件夹");
            directoryChooser.setInitialDirectory(new File(System.getProperty("user.home") + "/Downloads")
            );
            File directory = directoryChooser.showDialog(stage);
            if (null == directory) {
                System.out.println("用户取消选中文件夹");
                return;
            }
            //读取MNIST数据集
            try {
                readMnist(directory);
                warning.setText("读取成功");
            } catch (Exception ex) {
                warning.setText(ex.getMessage());
            }
        });
        //pic number input
        Text picNumberLabel = new Text("图");
        TextField picNumber = new TextField();
        pic = new SimpleStringProperty();
        pic.bind(picNumber.textProperty());
        //get pic button
        Button view = new Button("查看");
        view.setOnMouseClicked(event -> {
            if (null == trainImgReader) {
                warning.setText("请先选择有效的MNIST数据集文件夹");
                return;
            }
            String numberString = pic.get();
            int i = NumberUtils.toInt(numberString, -1);
            if (i < 0) {
                warning.setText("输入有误，应为有效正整数，请重新输入");
                return;
            }
            try {
                showPic(i);
                warning.setText("");
            } catch (Exception e) {
                warning.setText(e.getMessage());
            }
        });
        Button prevBtn = new Button("前一张");
        prevBtn.setOnMouseClicked(event -> {
            String numberString = pic.get();
            int i = NumberUtils.toInt(numberString, -1);
            --i;
            picNumber.setText(i + "");
            showPic(i);
        });
        Button nextBtn = new Button("后一张");
        nextBtn.setOnMouseClicked(event -> {
            String numberString = pic.get();
            int i = NumberUtils.toInt(numberString, -1);
            ++i;
            picNumber.setText(i + "");
            showPic(i);
        });
        FlowPane bottomControl = new FlowPane();
        bottomControl.getChildren().addAll(dirChooser, picNumberLabel, picNumber, label, view, warning, prevBtn, nextBtn);
        bottomControl.setPrefHeight(50);
        BorderPane root = new BorderPane();
        imageView = new ImageView();
        root.setCenter(imageView);
        root.setBottom(bottomControl);
        Scene scene = new Scene(root, DEFAULT_WIDTH, DEFAULT_HEIGHT + 50);
        stage.setScene(scene);
        stage.setTitle("Example5-Derivative");
        stage.setResizable(false);
        stage.show();
    }
}
